import requests
import json
from pathlib import Path

BASE_URL = "http://localhost:5000"

def upload_file(file_path, title, source):
    """上传文件到 RAG 系统"""
    url = f"{BASE_URL}/api/document"
    
    # 准备文件和其他表单数据
    files = {
        'file': (
            Path(file_path).name,
            open(file_path, 'rb'),
            get_content_type(file_path)
        )
    }
    data = {
        'title': title,
        'source': source
    }
    
    # 发送请求
    response = requests.post(url, files=files, data=data)
    
    # 打印结果
    print(f"\n上传文件: {file_path}")
    print(f"状态码: {response.status_code}")
    print("响应内容:")
    print(json.dumps(response.json(), ensure_ascii=False, indent=2))
    
    return response.json()

def get_content_type(file_path):
    """根据文件扩展名获取 Content-Type"""
    ext = Path(file_path).suffix.lower()
    content_types = {
        '.md': 'text/markdown',
        '.pdf': 'application/pdf',
        '.docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
        '.txt': 'text/plain'
    }
    return content_types.get(ext, 'application/octet-stream')

def main():
    # 上传 Markdown 文件
    upload_file(
        'test.md',
        '讯飞星火大模型介绍',
        '本地文档'
    )
    
    # # 上传 PDF 文件（取消注释并提供正确的文件路径）
    # upload_file(
    #     'test.pdf',
    #     'PDF测试文档',
    #     '本地文档'
    # )
    
    # # 上传 Word 文件（取消注释并提供正确的文件路径）
    # upload_file(
    #     'test.docx',
    #     'Word测试文档',
    #     '本地文档'
    # )

if __name__ == '__main__':
    main() 